# -*- coding: UTF-8 -*-
# @Project -> File   ：AIserving -> output_format
# @Author ：luoyang
# @Date   ：2021/3/2 19:32


class NerOutputFormat(object):

    def __init__(self):
        pass

    def change_json(self,ner_list):
        res = {'ner_result': []}
        for item in ner_list:
            new_item = {}
            new_item['sentence'] = item['sentence']
            if 'labels' in item:
                for label in item['labels']:
                    if label['entity'] not in new_item:
                        new_item[label['entity']] = []
                    new_item[label['entity']].append((label['start'], label['end'], label['value']))
            res['ner_result'].append(new_item)
        return res

    def parse_output(self, model_res):

        """
        对模型返回的结果进行解析，组装成最终返回结果
        :param model_res: 模型返回的结果，见pipeline.predict方法的返回结构
        :return:
            res: dict，形如：
                {
                    "ner_result":
                          [
                            {
                              "sentence": "收件地址：徐州市新城区新安路5号徐州市公共资源交易中心，收件联系人：周悦悦，收件联系方式：18811939611。",
                              "LOC": [[5, 27, "徐州市新城区新安路5号徐州市公共资源交易中心"]],
                              "PER": [[34, 37, "周悦悦"]]},
                            {
                              "sentence": "投标人采用委托人员送达方式递交U盘的，必须派法定代表人或其授权代理人携加盖公章的单位介绍信（或授权书）、本人身份证明，将U盘送达以下地点，送达地点：徐州市公共资源交易中心一楼北门口（徐州市新城区新安路5号，市规划馆南侧），招标代理机构联系人：周悦悦，电话：18811939611，送达指定地点前应主动与联系人周悦悦联系。",
                              "LOC": [[74, 90, "徐州市公共资源交易中心一楼北门口"], [91, 102, "徐州市新城区新安路5号"], [103, 109, "市规划馆南侧"]],
                              "ORG": [[111, 117, "招标代理机构"]],
                              "PER": [[121, 124, "周悦悦"], [154, 157, "周悦悦"]]
                            }
                          ]
                }
        注意：
            1、会把模型预测结果中，长度小于等于1的实体删除
        """

        # res = {
        #     'ner_result': []
        # }
        #
        # for every_res in model_res:
        #     sub_res = {'sentence': every_res['sentence']}
        #     for every in every_res['labels']:
        #         inds, inde, text, label = every
        #         if inde - inds > 1:
        #             if sub_res.get(label, 0):
        #                 sub_res[label].append((inds, inde, text))
        #             else:
        #                 sub_res[label] = [(inds, inde, text)]
        #     res['ner_result'].append(sub_res)

        return self.change_json(model_res)
